{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "ade0633c",
   "metadata": {},
   "source": [
    "QLearning使用表格估计Q函数,不便于扩展.\n",
    "\n",
    "所以使用神经网络估计Q函数."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "91162df8",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAASAAAADMCAYAAADTcn7NAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjYuMywgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/P9b71AAAACXBIWXMAAA9hAAAPYQGoP6dpAAAUYklEQVR4nO3dfWxT570H8K/t2A55OQ6BxSYiGUxDpREv2wIEr3+sd3hkXVSNNbraJtRlFaIqc1BpKqRFauml6r2pmHS7sUH6x+6gk25HlUlsakSL0gBBu3WhuI0UUohaiSoRYLvAcmwCfonP7/7BcqhLaBMS/Njx9yMdqed5Htu/c4K/PT6Pz7FFRARERApYVRdARIWLAUREyjCAiEgZBhARKcMAIiJlGEBEpAwDiIiUYQARkTIMICJShgFERMooC6B9+/ZhyZIlKC4uRkNDA06fPq2qFCJSREkAvfHGG2hra8MLL7yADz74AKtXr0ZjYyMikYiKcohIEYuKi1EbGhqwdu1a/OEPfwAAGIaBmpoabN++Hb/+9a+zXQ4RKVKU7RdMJpMIBoNob28326xWK3w+HwKBwKSPSSQSSCQS5rphGLh27RoWLFgAi8Vy32smoukREcRiMVRXV8NqvfsHrawH0JUrV5BOp+F2uzPa3W43zp8/P+ljOjo6sHv37myUR0SzaGRkBIsXL75rf9YD6F60t7ejra3NXNd1HbW1tRgZGYGmaQorI6LJRKNR1NTUoLy8/EvHZT2AFi5cCJvNhnA4nNEeDofh8XgmfYzT6YTT6byjXdM0BhBRDvuqUyRZnwVzOByor69Hb2+v2WYYBnp7e+H1erNdDhEppOQjWFtbG1paWrBmzRqsW7cOv/3tbzE2NoYnnnhCRTlEpIiSAPrpT3+Kzz77DLt27UIoFMK3vvUtvP3223ecmCaiuU3J94BmKhqNwuVyQdd1ngMiykFTfY/yWjAiUoYBRETKMICISBkGEBEpwwAiImUYQESkDAOIiJRhABGRMgwgIlKGAUREyjCAiEgZBhARKcMAIiJlGEBEpAwDiIiUYQARkTIMICJShgFERMowgIhIGQYQESnDACIiZRhARKQMA4iIlGEAEZEyDCAiUoYBRETKMICISBkGEBEpM+0AOnnyJB599FFUV1fDYrHgb3/7W0a/iGDXrl1YtGgR5s2bB5/Ph48//jhjzLVr17B582ZomoaKigps2bIF169fn9GGEFH+mXYAjY2NYfXq1di3b9+k/Xv27MHevXvx6quv4tSpUygtLUVjYyPi8bg5ZvPmzRgcHERPTw+6u7tx8uRJPPnkk/e+FUSUn2QGAMjhw4fNdcMwxOPxyG9+8xuzbXR0VJxOp/zlL38REZGPPvpIAMj7779vjnnrrbfEYrHIxYsXp/S6uq4LANF1fSblE9F9MtX36KyeA7pw4QJCoRB8Pp/Z5nK50NDQgEAgAAAIBAKoqKjAmjVrzDE+nw9WqxWnTp2a9HkTiQSi0WjGQkT5b1YDKBQKAQDcbndGu9vtNvtCoRCqqqoy+ouKilBZWWmO+aKOjg64XC5zqampmc2yiUiRvJgFa29vh67r5jIyMqK6JCKaBbMaQB6PBwAQDocz2sPhsNnn8XgQiUQy+sfHx3Ht2jVzzBc5nU5ompaxEFH+m9UAWrp0KTweD3p7e822aDSKU6dOwev1AgC8Xi9GR0cRDAbNMceOHYNhGGhoaJjNcogoxxVN9wHXr1/HJ598Yq5fuHAB/f39qKysRG1tLXbs2IGXXnoJy5Ytw9KlS/H888+juroamzZtAgA8+OCD+OEPf4itW7fi1VdfRSqVQmtrK372s5+hurp61jaMiPLAdKfXjh8/LgDuWFpaWkTk1lT8888/L263W5xOp2zYsEGGhoYynuPq1avy85//XMrKykTTNHniiSckFovN+hQfEakx1feoRUREYf7dk2g0CpfLBV3XeT6IKAdN9T2aF7NgRDQ3MYCISBkGEBEpwwAiImUYQESkDAOIiJRhABGRMgwgIlKGAUREyjCAiEgZBhARKcMAIiJlGEBEpAwDiIiUYQARkTIMICJShgFERMowgIhIGQYQESnDACIiZab9szxEsymdSkAfHoAY6X+1WKAtfhD2eeVK66LsYACRUuPx67jQ92cYqfi/Wix44NFnGUAFgh/BKOcY4wnVJVCWMIAo5xjjSdUlUJYwgCjHCIzxlOoiKEsYQKSUxWqDze7MaEvd0BVVQ9nGACKlbI5iOMoqM9oS0YiiaijbGECklMVihcXKydhCNa0A6ujowNq1a1FeXo6qqips2rQJQ0NDGWPi8Tj8fj8WLFiAsrIyNDc3IxwOZ4wZHh5GU1MTSkpKUFVVhZ07d2J8fHzmW0P5x2KF1cYAKlTTCqC+vj74/X6899576OnpQSqVwsaNGzE2NmaOeeaZZ/Dmm2+iq6sLfX19uHTpEh577DGzP51Oo6mpCclkEu+++y5ee+01HDx4ELt27Zq9raK8YbFYYWEAFS6ZgUgkIgCkr69PRERGR0fFbrdLV1eXOebcuXMCQAKBgIiIHDlyRKxWq4RCIXNMZ2enaJomiURiSq+r67oAEF3XZ1I+5QDDMOTjo/vl9KtbzeXCiT+LYRiqS6MZmOp7dEbngHT91mxFZeWtk4jBYBCpVAo+n88cs3z5ctTW1iIQCAAAAoEAVq5cCbfbbY5pbGxENBrF4ODgpK+TSCQQjUYzFpobLBYLAEtGm2GkAIiSeii77jmADMPAjh078NBDD2HFihUAgFAoBIfDgYqKioyxbrcboVDIHPP58Jnon+ibTEdHB1wul7nU1NTca9mUB4zxFCAMoEJwzwHk9/tx9uxZHDp0aDbrmVR7ezt0XTeXkZGR+/6apI6kxyEMoIJwT2f/Wltb0d3djZMnT2Lx4sVmu8fjQTKZxOjoaMZRUDgchsfjMcecPn064/kmZskmxnyR0+mE0+mctI/mHiPNI6BCMa0jIBFBa2srDh8+jGPHjmHp0qUZ/fX19bDb7ejt7TXbhoaGMDw8DK/XCwDwer0YGBhAJHL7y2Y9PT3QNA11dXUz2RbKU05tYcZ6Inb1VgjRnDetIyC/34/XX38df//731FeXm6es3G5XJg3bx5cLhe2bNmCtrY2VFZWQtM0bN++HV6vF+vXrwcAbNy4EXV1dXj88cexZ88ehEIhPPfcc/D7/TzKKVD20vkZ6+nkjc/dH4jmsmkFUGdnJwDg4Ycfzmg/cOAAfvnLXwIAXnnlFVitVjQ3NyORSKCxsRH79+83x9psNnR3d2Pbtm3wer0oLS1FS0sLXnzxxZltCeUta5FDdQmkiEXy8GxfNBqFy+WCruvQNE11OTRDVz4+hQvH/sdcL5pXjhX//h+8KVkem+p7lNeCkXI2HgEVLAYQKXfHpRgCzoIVCAYQ5SDhLFiBYABRzhExIGneHaEQMIAo9wiPgAoFA4iUu3VB6m1iGLwxfYFgAJFyjtJK2Jwl5roxnkCct2UtCAwgUs5a5LjztqyGoaYYyioGEClnsRXBYuE/xULEvzopZ7EVwWLlP8VCxL86KWe1TnYEJLwnUAFgAJF6Fsut5XOMNK+GLwQMIMpJnIYvDAwgykkMoMLAAKKcxAAqDAwgUs5iscBqL85oS92MKaqGsokBRMpZrDbMq8j8qab46OQ/0URzCwOIcoAFVhtvSlaIGECkngWwFNlVV0EKMIAoB1h4Y/oCxQCinGC1ffEISMDfh5/7GECknGWSb0KLYUB4RfycxwCinCTGOAOoADCAKCeJkYYIrweb66b1y6hE90pEEIvFYNzlqCYRj2esj6eS0Ef/CZsjMel4q9WK8vLyO27nSvmFAURZEY/H0dTUhAsXLkzav2FVNbY9ugrWf92WY/Szy9jq+zdcvjY26fiamhq88847KC0tvW810/3HAKKsCYfDuHjx4qR9n1RYoacexqeJNRBY8HVHEP+8+hkuXvrnpOPtdjvvFzQHTOscUGdnJ1atWgVN06BpGrxeL9566y2zPx6Pw+/3Y8GCBSgrK0NzczPC4XDGcwwPD6OpqQklJSWoqqrCzp07MT7O34AqdHrCgQ+jPoST30AkuRT9MR/iRpnqsug+m1YALV68GC+//DKCwSDOnDmD73//+/jxj3+MwcFBAMAzzzyDN998E11dXejr68OlS5fw2GOPmY9Pp9NoampCMpnEu+++i9deew0HDx7Erl27ZnerKO+MJay4kb79cSqeLkVK+OXEOU9maP78+fLHP/5RRkdHxW63S1dXl9l37tw5ASCBQEBERI4cOSJWq1VCoZA5prOzUzRNk0QiMeXX1HVdAIiu6zMtn7Lkxo0bsmzZsolvF96xLKmukr3/+b+y+6VTsvulU7L3v96Qry9y3338kiUSi8VUbxbdxVTfo/d8DiidTqOrqwtjY2Pwer0IBoNIpVLw+XzmmOXLl6O2thaBQADr169HIBDAypUr4XbfvvK5sbER27Ztw+DgIL797W9Pq4bz58+jrIyH6fkgHo8jmbz7PX6u6Tr+7/h/48r4NyBiQZX9AvSYftfxqVQK58+fR0lJyV3HkDrXr1+f0rhpB9DAwAC8Xi/i8TjKyspw+PBh1NXVob+/Hw6HAxUVFRnj3W43QqFbt1YIhUIZ4TPRP9F3N4lEAonE7enYaDQKANB1neeP8kQikbjrFDwARMcSeOOdIIDglJ7PMAzouv6loUbqjI1NPnv5RdMOoAceeAD9/f3QdR1//etf0dLSgr6+vmkXOB0dHR3YvXv3He0NDQ3QNO2+vjbNjps3b6K4uPirB06R0+lEQ0MDj4Bz1MRBwleZ9jehHQ4HvvnNb6K+vh4dHR1YvXo1fve738Hj8SCZTGJ0dDRjfDgchsfjAQB4PJ47ZsUm1ifGTKa9vR26rpvLyMjIdMsmohw040sxDMNAIpFAfX097HY7ent7zb6hoSEMDw/D6/UCALxeLwYGBhCJ3P7d756eHmiahrq6uru+htPpNKf+JxYiyn/T+gjW3t6ORx55BLW1tYjFYnj99ddx4sQJHD16FC6XC1u2bEFbWxsqKyuhaRq2b98Or9eL9evXAwA2btyIuro6PP7449izZw9CoRCee+45+P1+OJ3O+7KBRJS7phVAkUgEv/jFL3D58mW4XC6sWrUKR48exQ9+8AMAwCuvvAKr1Yrm5mYkEgk0NjZi//795uNtNhu6u7uxbds2eL1elJaWoqWlBS+++OLsbhXlpPLy8lk7ei0vL5+V5yG1LCL59332aDQKl8sFXdf5cSxPGIaBcDg8a7OWNpsNHo8HVv6mfE6a6nuU14JRVlitVixatEh1GZRj+L8PIlKGAUREyjCAiEgZBhARKcMAIiJlGEBEpAwDiIiUYQARkTIMICJShgFERMowgIhIGQYQESnDACIiZRhARKQMA4iIlGEAEZEyDCAiUoYBRETKMICISBkGEBEpwwAiImUYQESkDAOIiJRhABGRMgwgIlKGAUREyjCAiEgZBhARKcMAIiJlGEBEpEyR6gLuhYgAAKLRqOJKiGgyE+/Niffq3eRlAF29ehUAUFNTo7gSIvoysVgMLpfrrv15GUCVlZUAgOHh4S/dOMoUjUZRU1ODkZERaJqmupy8wH12b0QEsVgM1dXVXzouLwPIar116srlcvEfxT3QNI37bZq4z6ZvKgcHPAlNRMowgIhImbwMIKfTiRdeeAFOp1N1KXmF+236uM/uL4t81TwZEdF9kpdHQEQ0NzCAiEgZBhARKcMAIiJl8jKA9u3bhyVLlqC4uBgNDQ04ffq06pKU6ejowNq1a1FeXo6qqips2rQJQ0NDGWPi8Tj8fj8WLFiAsrIyNDc3IxwOZ4wZHh5GU1MTSkpKUFVVhZ07d2J8fDybm6LMyy+/DIvFgh07dpht3GdZInnm0KFD4nA45E9/+pMMDg7K1q1bpaKiQsLhsOrSlGhsbJQDBw7I2bNnpb+/X370ox9JbW2tXL9+3Rzz1FNPSU1NjfT29sqZM2dk/fr18t3vftfsHx8flxUrVojP55MPP/xQjhw5IgsXLpT29nYVm5RVp0+fliVLlsiqVavk6aefNtu5z7Ij7wJo3bp14vf7zfV0Oi3V1dXS0dGhsKrcEYlEBID09fWJiMjo6KjY7Xbp6uoyx5w7d04ASCAQEBGRI0eOiNVqlVAoZI7p7OwUTdMkkUhkdwOyKBaLybJly6Snp0e+973vmQHEfZY9efURLJlMIhgMwufzmW1WqxU+nw+BQEBhZblD13UAty/YDQaDSKVSGfts+fLlqK2tNfdZIBDAypUr4Xa7zTGNjY2IRqMYHBzMYvXZ5ff70dTUlLFvAO6zbMqri1GvXLmCdDqd8UcHALfbjfPnzyuqKncYhoEdO3bgoYcewooVKwAAoVAIDocDFRUVGWPdbjdCoZA5ZrJ9OtE3Fx06dAgffPAB3n///Tv6uM+yJ68CiL6c3+/H2bNn8Y9//EN1KTltZGQETz/9NHp6elBcXKy6nIKWVx/BFi5cCJvNdsdsRDgchsfjUVRVbmhtbUV3dzeOHz+OxYsXm+0ejwfJZBKjo6MZ4z+/zzwez6T7dKJvrgkGg4hEIvjOd76DoqIiFBUVoa+vD3v37kVRURHcbjf3WZbkVQA5HA7U19ejt7fXbDMMA729vfB6vQorU0dE0NraisOHD+PYsWNYunRpRn99fT3sdnvGPhsaGsLw8LC5z7xeLwYGBhCJRMwxPT090DQNdXV12dmQLNqwYQMGBgbQ399vLmvWrMHmzZvN/+Y+yxLVZ8Gn69ChQ+J0OuXgwYPy0UcfyZNPPikVFRUZsxGFZNu2beJyueTEiRNy+fJlc7lx44Y55qmnnpLa2lo5duyYnDlzRrxer3i9XrN/Ykp548aN0t/fL2+//bZ87WtfK6gp5c/Pgolwn2VL3gWQiMjvf/97qa2tFYfDIevWrZP33ntPdUnKAJh0OXDggDnm5s2b8qtf/Urmz58vJSUl8pOf/EQuX76c8TyffvqpPPLIIzJv3jxZuHChPPvss5JKpbK8Nep8MYC4z7KDt+MgImXy6hwQEc0tDCAiUoYBRETKMICISBkGEBEpwwAiImUYQESkDAOIiJRhABGRMgwgIlKGAUREyjCAiEiZ/wfksNbQHV4vFAAAAABJRU5ErkJggg==\n",
      "text/plain": [
       "<Figure size 300x300 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import gym\n",
    "\n",
    "\n",
    "#定义环境\n",
    "class MyWrapper(gym.Wrapper):\n",
    "\n",
    "    def __init__(self):\n",
    "        env = gym.make('CartPole-v1', render_mode='rgb_array')\n",
    "        super().__init__(env)\n",
    "        self.env = env\n",
    "        self.step_n = 0\n",
    "\n",
    "    def reset(self):\n",
    "        state, _ = self.env.reset()\n",
    "        self.step_n = 0\n",
    "        return state\n",
    "\n",
    "    def step(self, action):\n",
    "        state, reward, terminated, truncated, info = self.env.step(action)\n",
    "        over = terminated or truncated\n",
    "\n",
    "        #限制最大步数\n",
    "        self.step_n += 1\n",
    "        if self.step_n >= 200:\n",
    "            over = True\n",
    "        \n",
    "        #没坚持到最后,扣分\n",
    "        if over and self.step_n < 200:\n",
    "            reward = -1000\n",
    "\n",
    "        return state, reward, over\n",
    "\n",
    "    #打印游戏图像\n",
    "    def show(self):\n",
    "        from matplotlib import pyplot as plt\n",
    "        plt.figure(figsize=(3, 3))\n",
    "        plt.imshow(self.env.render())\n",
    "        plt.show()\n",
    "\n",
    "\n",
    "env = MyWrapper()\n",
    "\n",
    "env.reset()\n",
    "\n",
    "env.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "ecfbe912",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Sequential(\n",
       "  (0): Linear(in_features=4, out_features=64, bias=True)\n",
       "  (1): ReLU()\n",
       "  (2): Linear(in_features=64, out_features=64, bias=True)\n",
       "  (3): ReLU()\n",
       "  (4): Linear(in_features=64, out_features=2, bias=True)\n",
       ")"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import torch\n",
    "\n",
    "#定义模型,评估状态下每个动作的价值\n",
    "model = torch.nn.Sequential(\n",
    "    torch.nn.Linear(4, 64),\n",
    "    torch.nn.ReLU(),\n",
    "    torch.nn.Linear(64, 64),\n",
    "    torch.nn.ReLU(),\n",
    "    torch.nn.Linear(64, 2),\n",
    ")\n",
    "\n",
    "model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "84cbf0ff",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "-991.0"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from IPython import display\n",
    "import random\n",
    "\n",
    "\n",
    "#玩一局游戏并记录数据\n",
    "def play(show=False):\n",
    "    data = []\n",
    "    reward_sum = 0\n",
    "\n",
    "    state = env.reset()\n",
    "    over = False\n",
    "    while not over:\n",
    "        action = model(torch.FloatTensor(state).reshape(1, 4)).argmax().item()\n",
    "        if random.random() < 0.1:\n",
    "            action = env.action_space.sample()\n",
    "\n",
    "        next_state, reward, over = env.step(action)\n",
    "\n",
    "        data.append((state, action, reward, next_state, over))\n",
    "        reward_sum += reward\n",
    "\n",
    "        state = next_state\n",
    "\n",
    "        if show:\n",
    "            display.clear_output(wait=True)\n",
    "            env.show()\n",
    "\n",
    "    return data, reward_sum\n",
    "\n",
    "\n",
    "play()[-1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "95ce610e",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/tmp/ipykernel_23708/2282743005.py:27: UserWarning: Creating a tensor from a list of numpy.ndarrays is extremely slow. Please consider converting the list to a single numpy.ndarray with numpy.array() before converting to a tensor. (Triggered internally at  ../torch/csrc/utils/tensor_new.cpp:201.)\n",
      "  state = torch.FloatTensor([i[0] for i in data]).reshape(-1, 4)\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "(208,\n",
       " (array([0.00496458, 0.03820252, 0.0385537 , 0.00248376], dtype=float32),\n",
       "  1,\n",
       "  1.0,\n",
       "  array([ 0.00572863,  0.23275095,  0.03860337, -0.27779007], dtype=float32),\n",
       "  False))"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#数据池\n",
    "class Pool:\n",
    "\n",
    "    def __init__(self):\n",
    "        self.pool = []\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.pool)\n",
    "\n",
    "    def __getitem__(self, i):\n",
    "        return self.pool[i]\n",
    "\n",
    "    #更新动作池\n",
    "    def update(self):\n",
    "        #每次更新不少于N条新数据\n",
    "        old_len = len(self.pool)\n",
    "        while len(pool) - old_len < 200:\n",
    "            self.pool.extend(play()[0])\n",
    "\n",
    "        #只保留最新的N条数据\n",
    "        self.pool = self.pool[-2_0000:]\n",
    "\n",
    "    #获取一批数据样本\n",
    "    def sample(self):\n",
    "        data = random.sample(self.pool, 64)\n",
    "\n",
    "        state = torch.FloatTensor([i[0] for i in data]).reshape(-1, 4)\n",
    "        action = torch.LongTensor([i[1] for i in data]).reshape(-1, 1)\n",
    "        reward = torch.FloatTensor([i[2] for i in data]).reshape(-1, 1)\n",
    "        next_state = torch.FloatTensor([i[3] for i in data]).reshape(-1, 4)\n",
    "        over = torch.LongTensor([i[4] for i in data]).reshape(-1, 1)\n",
    "\n",
    "        return state, action, reward, next_state, over\n",
    "\n",
    "\n",
    "pool = Pool()\n",
    "pool.update()\n",
    "pool.sample()\n",
    "\n",
    "len(pool), pool[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "018c4b1e",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0 415 -877.85\n",
      "100 20000 -813.6\n",
      "200 20000 200.0\n",
      "300 20000 200.0\n",
      "400 20000 200.0\n",
      "500 20000 200.0\n",
      "600 20000 200.0\n",
      "700 20000 -96.9\n",
      "800 20000 200.0\n",
      "900 20000 200.0\n"
     ]
    }
   ],
   "source": [
    "#训练\n",
    "def train():\n",
    "    model.train()\n",
    "    optimizer = torch.optim.Adam(model.parameters(), lr=2e-4)\n",
    "    loss_fn = torch.nn.MSELoss()\n",
    "\n",
    "    #共更新N轮数据\n",
    "    for epoch in range(1000):\n",
    "        pool.update()\n",
    "\n",
    "        #每次更新数据后,训练N次\n",
    "        for i in range(200):\n",
    "\n",
    "            #采样N条数据\n",
    "            state, action, reward, next_state, over = pool.sample()\n",
    "\n",
    "            #计算value\n",
    "            value = model(state).gather(dim=1, index=action)\n",
    "\n",
    "            #计算target\n",
    "            with torch.no_grad():\n",
    "                target = model(next_state)\n",
    "            target = target.max(dim=1)[0].reshape(-1, 1)\n",
    "            target = target * 0.99 * (1 - over) + reward\n",
    "\n",
    "            loss = loss_fn(value, target)\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "            optimizer.zero_grad()\n",
    "\n",
    "        if epoch % 100 == 0:\n",
    "            test_result = sum([play()[-1] for _ in range(20)]) / 20\n",
    "            print(epoch, len(pool), test_result)\n",
    "\n",
    "\n",
    "train()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "a101d0e6",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAASAAAADMCAYAAADTcn7NAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjYuMywgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/P9b71AAAACXBIWXMAAA9hAAAPYQGoP6dpAAAVB0lEQVR4nO3dfWxT570H8K9fYuf1OA00djPiwbTe0oiXdQGCW121d81Iu2gabf7YJsTSCtHbzEGlqZAWiVKBqhsu+6NdVxr+mAb80zFlEq2aC62yAEG9NW/pooUA2brbLRFgp4HFTkLit/O7f7CcYpJATIyfGL4f6Uj4eR7bv3PgfDl+js+xSUQEREQKmFUXQET3LwYQESnDACIiZRhARKQMA4iIlGEAEZEyDCAiUoYBRETKMICISBkGEBEpoyyAdu/ejYULFyI7OxsVFRU4deqUqlKISBElAfT73/8eDQ0NeOONN/D5559j+fLlqKqqwsDAgIpyiEgRk4qLUSsqKrBy5Uq8++67AABd11FaWopNmzbhF7/4RbrLISJFrOl+w0gkgs7OTjQ2NhptZrMZlZWV8Pl8Uz4nHA4jHA4bj3Vdx9WrVzFv3jyYTKa7XjMRJUdEMDw8jJKSEpjN03/QSnsADQ4OIh6Pw+l0JrQ7nU5cuHBhyuc0NTVh+/bt6SiPiFKov78fCxYsmLY/7QF0JxobG9HQ0GA8DgaDcLvd6O/vh6ZpCisjoqmEQiGUlpaioKDgluPSHkDz58+HxWJBIBBIaA8EAnC5XFM+x263w263T2rXNI0BRDSH3W6KJO1nwWw2G8rLy9He3m606bqO9vZ2eDyedJdDRAop+QjW0NCA2tparFixAqtWrcLbb7+N0dFRvPjiiyrKISJFlATQj3/8Y3z11VfYtm0b/H4/vvOd7+Djjz+eNDFNRPc2Jd8Dmq1QKASHw4FgMMg5IKI5aKb7KK8FIyJlGEBEpAwDiIiUYQARkTIMICJShgFERMowgIhIGQYQESnDACIiZRhARKQMA4iIlGEAEZEyDCAiUoYBRETKMICISBkGEBEpwwAiImUYQESkDAOIiJRhABGRMgwgIlKGAUREyjCAiEgZBhARKcMAIiJlGEBEpAwDiIiUYQARkTJJB9Dx48fxwx/+ECUlJTCZTPjggw8S+kUE27Ztw0MPPYScnBxUVlbir3/9a8KYq1evYt26ddA0DYWFhdiwYQNGRkZmtSJElHmSDqDR0VEsX74cu3fvnrJ/165deOedd7Bnzx6cPHkSeXl5qKqqwvj4uDFm3bp16OnpQVtbG1pbW3H8+HG89NJLd74WRJSZZBYAyMGDB43Huq6Ly+WSX/7yl0bb0NCQ2O12+d3vficiIufOnRMAcvr0aWPM4cOHxWQyycWLF2f0vsFgUABIMBicTflEdJfMdB9N6RzQl19+Cb/fj8rKSqPN4XCgoqICPp8PAODz+VBYWIgVK1YYYyorK2E2m3Hy5MkpXzccDiMUCiUsRJT5UhpAfr8fAOB0OhPanU6n0ef3+1FcXJzQb7VaUVRUZIy5WVNTExwOh7GUlpamsmwiUiQjzoI1NjYiGAwaS39/v+qSiCgFUhpALpcLABAIBBLaA4GA0edyuTAwMJDQH4vFcPXqVWPMzex2OzRNS1iIKPOlNIAWLVoEl8uF9vZ2oy0UCuHkyZPweDwAAI/Hg6GhIXR2dhpjjhw5Al3XUVFRkcpyiGiOsyb7hJGREXzxxRfG4y+//BJdXV0oKiqC2+3G5s2b8eabb+Lhhx/GokWL8Prrr6OkpARr164FADz66KN45plnsHHjRuzZswfRaBT19fX4yU9+gpKSkpStGBFlgGRPrx09elQATFpqa2tF5Pqp+Ndff12cTqfY7XZ5+umnpbe3N+E1rly5Ij/96U8lPz9fNE2TF198UYaHh1N+io+I1JjpPmoSEVGYf3ckFArB4XAgGAxyPohoDprpPpoRZ8GI6N7EACIiZRhARKQMA4iIlGEAEZEyDCAiUoYBRETKMICISBkGEBEpwwAiImUYQESkDAOIiJRhABGRMgwgIlKGAUREyjCAiEgZBhARKcMAIiJlGEBEpAwDiIiUSfpneSiRHosgOjaM2PgIYuFRAID2jcUwmZjtRLfDAJqlwb/4cPH0h9DjMYgeQ7ZWjEefa4Qly666NKI5jwE0S3o0gtj4iPE4HgtDj4UZQEQzwM8Js2S22QGYjMexsWFErwXVFUSUQRhAs5RbtABmy9cHknosCj0WUVgRUeZgAM2SxZ4LmBM3Ywb+2CyREgygWTJbbbjxIxgAxCNjaoohyjAMoJQTRK+FVBdBlBGSCqCmpiasXLkSBQUFKC4uxtq1a9Hb25swZnx8HF6vF/PmzUN+fj5qamoQCAQSxvT19aG6uhq5ubkoLi7Gli1bEIvFZr82CpitNmTlOhLarg3+Q1E1RJklqQDq6OiA1+vFiRMn0NbWhmg0ijVr1mB0dNQY8+qrr+Kjjz5CS0sLOjo6cOnSJTz//PNGfzweR3V1NSKRCD777DPs378f+/btw7Zt21K3VmlksdphyytMaBM9rqYYokwjszAwMCAApKOjQ0REhoaGJCsrS1paWowx58+fFwDi8/lEROTQoUNiNpvF7/cbY5qbm0XTNAmHwzN632AwKAAkGAzOpvyU0PW4/OXwu3Jqz0Zj+b9j+0XXddWlESkz0310VnNAweD177sUFRUBADo7OxGNRlFZWWmMWbx4MdxuN3w+HwDA5/Nh6dKlcDqdxpiqqiqEQiH09PRM+T7hcBihUChhmTtMMJktCS3xyBgguqJ6iDLHHQeQruvYvHkznnjiCSxZsgQA4Pf7YbPZUFhYmDDW6XTC7/cbY24Mn4n+ib6pNDU1weFwGEtpaemdlp1yJpMJpptOw8fDo/wYRjQDdxxAXq8XZ8+exYEDB1JZz5QaGxsRDAaNpb+//66/ZzJy5iUG4tiQn19GJJqBO7oWrL6+Hq2trTh+/DgWLFhgtLtcLkQiEQwNDSUcBQUCAbhcLmPMqVOnEl5v4izZxJib2e122O1z99qqm8+CiR7nlxGJZiCpIyARQX19PQ4ePIgjR45g0aJFCf3l5eXIyspCe3u70dbb24u+vj54PB4AgMfjQXd3NwYGBowxbW1t0DQNZWVls1kXZaz2vMQGAeeAiGYgqSMgr9eL999/Hx9++CEKCgqMORuHw4GcnBw4HA5s2LABDQ0NKCoqgqZp2LRpEzweD1avXg0AWLNmDcrKyrB+/Xrs2rULfr8fW7duhdfrndNHObdy8yS0SBzxyPikIyMiSpRUADU3NwMAnnrqqYT2vXv34oUXXgAAvPXWWzCbzaipqUE4HEZVVRXee+89Y6zFYkFrayvq6urg8XiQl5eH2tpa7NixY3ZropDJlHgphh6LInJtCNmFzmmeQUQAYJIMnKwIhUJwOBwIBoPQNE11ORgb8uPCB/9t3BERABY99QLmP/K4wqqI1JnpPsprwVLAas/910WpRJQMBlAKmK12mCyJn2b1eIxnwohugwGUCiYTbr4lx423aSWiqTGAUsAETLoHdGTkippiiDIIAygFTGYLcoq+kdA29s+pLyshoq8xgFLBZIY1O+/244goAQMoRSxZ2QmPRY/zglSi22AApYBpikloPRaGHo+qKYgoQzCAUsR88yT06BBiY8OKqiHKDAygFMmd7064JkyPjvOWHES3wQBKEWt23r++D0REM8UASpGpfh9Mj4bVFEOUIRhAKXLzsY+IJFycSkSTMYBSxJyVDWt2/tcNomPsn5fUFUSUARhAKWKx5SArN/G2A7ExXg9GdCsMoBQxW6yTv4wI4RXxRLfAAEoVk3nSrVmvT0IzgIimwwBKoZtvzRobH4HovDk90XQYQCmUU7Qg4fG1K/2QeExRNURzHwMohbLyChMeX78YlR/BiKbDAEqhSb8PBoBz0ETTYwCl0KTfB4vHoMfGFVVDNPcxgFLEZDJN/n0wPYZ4ZExRRURzHwMohexaMSxZOcbjeHgM48GvFFZENLcl9cuo9zsRwejoKGKxqc9sRSNxiPnGTBeMjoRgGhqa9jULCgpgsVim7Se6lzGAklRXV4ejR49O2ZeXbcXb//nvmFfw9VHQaw0N+N9zl6ccn5ubiz/+8Y9wu913pVaiuY4BlKTBwUFcvHhxyr78HBtGIzm4NPIkwnouvplzFtHxT6cdn5eXN+3RFNH9IKk5oObmZixbtgyapkHTNHg8Hhw+fNjoHx8fh9frxbx585Cfn4+amhoEAoGE1+jr60N1dTVyc3NRXFyMLVu23DM7YVys+PPIU7gY/jcMRkvx5+H/QI7jW6rLIpqzkgqgBQsWYOfOnejs7MSZM2fwve99Dz/60Y/Q09MDAHj11Vfx0UcfoaWlBR0dHbh06RKef/554/nxeBzV1dWIRCL47LPPsH//fuzbtw/btm1L7VoposOM0bgDE3cHiko2Sh8qVVsU0Vwms/TAAw/Ib37zGxkaGpKsrCxpaWkx+s6fPy8AxOfziYjIoUOHxGw2i9/vN8Y0NzeLpmkSDodn/J7BYFAASDAYnG35SdF1XZ555hnB9a83T1osFqtsWL9Vdrx5Qra/eUre3vk/8l91NdOOz8vLk7/97W9pXQeidJjpPnrHc0DxeBwtLS0YHR2Fx+NBZ2cnotEoKisrjTGLFy+G2+2Gz+fD6tWr4fP5sHTpUjidTmNMVVUV6urq0NPTg8ceeyypGi5cuID8/PzbD0yhkZHp7/ETj8fQdXo/+vvPIiI5cNr+gcErU09AA4Cu6/jiiy8wPs4vK9K95Vb7yY2SDqDu7m54PB6Mj48jPz8fBw8eRFlZGbq6umCz2VBYWJgw3ul0wu+//jPFfr8/IXwm+if6phMOhxEOf31/5VAoBAAIBoNpnz+63ft1XugHLvTP6LVEBMPDwxi6xWl6okw0Ojqz2xEnHUCPPPIIurq6EAwG8Yc//AG1tbXo6OhIusBkNDU1Yfv27ZPaKyoqoGnaFM+4O0RkUsDOhsViwWOPPYZvfYsT1XRvmThIuJ2kvwlts9nw7W9/G+Xl5WhqasLy5cvxq1/9Ci6XC5FIZNL/5oFAAC6XCwDgcrkmnRWbeDwxZiqNjY0IBoPG0t8/syMMIprbZn0phq7rCIfDKC8vR1ZWFtrb242+3t5e9PX1wePxAAA8Hg+6u7sxMDBgjGlra4OmaSgrK5v2Pex2u3Hqf2IhosyX1EewxsZGPPvss3C73RgeHsb777+PY8eO4ZNPPoHD4cCGDRvQ0NCAoqIiaJqGTZs2wePxYPXq1QCANWvWoKysDOvXr8euXbvg9/uxdetWeL1e2O3227w7Ed1rkgqggYEB/OxnP8Ply5fhcDiwbNkyfPLJJ/j+978PAHjrrbdgNptRU1ODcDiMqqoqvPfee8bzLRYLWltbUVdXB4/Hg7y8PNTW1mLHjh2pXau7KDc3N2VHYHl5eTCbeT0w3b9MIpl3y6xQKASHw4FgMJj2SejBwcGUnTY3mUxwuVywWnlFDN1bZrqP8l9+EkwmEx588EHVZRDdM3j8T0TKMICISBkGEBEpwwAiImUYQESkDAOIiJRhABGRMgwgIlKGAUREyjCAiEgZBhARKcMAIiJlGEBEpAwDiIiUYQARkTIMICJShgFERMowgIhIGQYQESnDACIiZRhARKQMA4iIlGEAEZEyDCAiUoYBRETKMICISBkGEBEpwwAiImUYQESkDAOIiJSxqi7gTogIACAUCimuhIimMrFvTuyr08nIALpy5QoAoLS0VHElRHQrw8PDcDgc0/ZnZAAVFRUBAPr6+m65cpQoFAqhtLQU/f390DRNdTkZgdvszogIhoeHUVJScstxGRlAZvP1qSuHw8F/FHdA0zRutyRxmyVvJgcHnIQmImUYQESkTEYGkN1uxxtvvAG73a66lIzC7ZY8brO7yyS3O09GRHSXZOQREBHdGxhARKQMA4iIlGEAEZEyGRlAu3fvxsKFC5GdnY2KigqcOnVKdUnKNDU1YeXKlSgoKEBxcTHWrl2L3t7ehDHj4+Pwer2YN28e8vPzUVNTg0AgkDCmr68P1dXVyM3NRXFxMbZs2YJYLJbOVVFm586dMJlM2Lx5s9HGbZYmkmEOHDggNptNfvvb30pPT49s3LhRCgsLJRAIqC5NiaqqKtm7d6+cPXtWurq65Ac/+IG43W4ZGRkxxrz88stSWloq7e3tcubMGVm9erU8/vjjRn8sFpMlS5ZIZWWl/OlPf5JDhw7J/PnzpbGxUcUqpdWpU6dk4cKFsmzZMnnllVeMdm6z9Mi4AFq1apV4vV7jcTwel5KSEmlqalJY1dwxMDAgAKSjo0NERIaGhiQrK0taWlqMMefPnxcA4vP5RETk0KFDYjabxe/3G2Oam5tF0zQJh8PpXYE0Gh4elocfflja2trkySefNAKI2yx9MuojWCQSQWdnJyorK402s9mMyspK+Hw+hZXNHcFgEMDXF+x2dnYiGo0mbLPFixfD7XYb28zn82Hp0qVwOp3GmKqqKoRCIfT09KSx+vTyer2orq5O2DYAt1k6ZdTFqIODg4jH4wl/6QDgdDpx4cIFRVXNHbquY/PmzXjiiSewZMkSAIDf74fNZkNhYWHCWKfTCb/fb4yZaptO9N2LDhw4gM8//xynT5+e1Mdtlj4ZFUB0a16vF2fPnsWnn36qupQ5rb+/H6+88gra2tqQnZ2tupz7WkZ9BJs/fz4sFsuksxGBQAAul0tRVXNDfX09WltbcfToUSxYsMBod7lciEQiGBoaShh/4zZzuVxTbtOJvntNZ2cnBgYG8N3vfhdWqxVWqxUdHR145513YLVa4XQ6uc3SJKMCyGazoby8HO3t7Uabrutob2+Hx+NRWJk6IoL6+nocPHgQR44cwaJFixL6y8vLkZWVlbDNent70dfXZ2wzj8eD7u5uDAwMGGPa2tqgaRrKysrSsyJp9PTTT6O7uxtdXV3GsmLFCqxbt874M7dZmqieBU/WgQMHxG63y759++TcuXPy0ksvSWFhYcLZiPtJXV2dOBwOOXbsmFy+fNlYrl27Zox5+eWXxe12y5EjR+TMmTPi8XjE4/EY/ROnlNesWSNdXV3y8ccfy4MPPnhfnVK+8SyYCLdZumRcAImI/PrXvxa32y02m01WrVolJ06cUF2SMgCmXPbu3WuMGRsbk5///OfywAMPSG5urjz33HNy+fLlhNf5+9//Ls8++6zk5OTI/Pnz5bXXXpNoNJrmtVHn5gDiNksP3o6DiJTJqDkgIrq3MICISBkGEBEpwwAiImUYQESkDAOIiJRhABGRMgwgIlKGAUREyjCAiEgZBhARKcMAIiJl/h/hkAugUDBdSAAAAABJRU5ErkJggg==\n",
      "text/plain": [
       "<Figure size 300x300 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "200.0"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "play(True)[-1]"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:pt39]",
   "language": "python",
   "name": "conda-env-pt39-py"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
